import torch
from PIL import Image
from transformers import AutoModel, AutoTokenizer

import util

model = AutoModel.from_pretrained('hf-models/MiniCPM-Llama3-V-2_5', trust_remote_code=True, torch_dtype=torch.float16)
model = model.to(device='cuda')

tokenizer = AutoTokenizer.from_pretrained('hf-models/MiniCPM-Llama3-V-2_5', trust_remote_code=True)
model.eval()


def recognize(image_url, target_text):
    util.download_and_show_image(image_url)
    image = Image.open(util.image_path).convert('RGB')
    question = f'''
    你现在的任务是从图片中提取相关信息，提取目标包括：{target_text}，各提取目标使用|分隔，请自行拆分，不要随意识别不相关信息。
    要求：
    1. 返回结果使用json格式，每条数据包含一个key-value对，key值为我指定的关键信息，value值为所抽取的结果。
    2. 仅考虑单条结果情况。
    3. 如果认为OCR识别结果中没有关键信息key，则将value赋值为“未找到相关信息”即可。如果图片模糊，请不要随意猜测内容，将value赋值为“无法识别”即可。 
    4. 请只输出json格式的结果，不要包含其它多余文字。
    '''
    msgs = [{'role': 'user', 'content': question}]

    res = model.chat(
        image=image,
        msgs=msgs,
        tokenizer=tokenizer,
        sampling=True, # if sampling=False, beam_search will be used by default
        temperature=0.7,
        # system_prompt='' # pass system_prompt if needed
    )
    print(res)

    return res
